from diffusers import DiffusionPipeline
import torch
import gradio as gr
import random

pipeline = DiffusionPipeline.from_pretrained("hf-models/ddpm-cat-256", torch_dtype=torch.float16)
pipeline.to("cuda")

def predict(steps, seed):
    generator = torch.manual_seed(seed)
    for i in range(1,steps):
        yield pipeline(generator=generator, num_inference_steps=i).images[0]

random_seed = random.randint(0, 2147483647)
gr.Interface(
    predict,
    inputs=[
        gr.inputs.Slider(1, 100, label='Inference Steps', default=5, step=1),
        gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed, step=1),
    ],
    outputs=gr.Image(shape=[128,128], type="pil", elem_id="output_image"),
    css="#output_image{width: 256px}",
    title="Unconditional butterflies",
    description="图片生成器",
).queue().launch()

